import cv2
import numpy as np


def find_template_in_screenshot(template_path, screenshot_path, threshold=0.8):
    """
    使用模板匹配在屏幕截图中查找模板图片的位置
    
    Args:
        template_path (str): 模板图片路径
        screenshot_path (str): 屏幕截图路径
        threshold (float): 匹配阈值，默认为0.8
    
    Returns:
        list: 包含匹配位置的列表，每个元素为 (top_left, bottom_right) 坐标元组
    """
    
    # 读取模板图片和屏幕截图
    template = cv2.imread(template_path, cv2.IMREAD_UNCHANGED)
    screenshot = cv2.imread(screenshot_path, cv2.IMREAD_UNCHANGED)
    
    if template is None:
        raise FileNotFoundError(f"无法读取模板图片: {template_path}")
    
    if screenshot is None:
        raise FileNotFoundError(f"无法读取屏幕截图: {screenshot_path}")
    
    # 转换为灰度图进行匹配
    if len(template.shape) == 3:
        template_gray = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
    else:
        template_gray = template
    
    if len(screenshot.shape) == 3:
        screenshot_gray = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)
    else:
        screenshot_gray = screenshot
    
    # 执行模板匹配
    result = cv2.matchTemplate(screenshot_gray, template_gray, cv2.TM_CCOEFF_NORMED)
    
    # 查找匹配位置
    locations = np.where(result >= threshold)
    
    # 获取模板的尺寸
    template_h, template_w = template_gray.shape
    
    # 存储匹配位置
    positions = []
    for pt in zip(*locations[::-1]):
        top_left = pt
        bottom_right = (pt[0] + template_w, pt[1] + template_h)
        positions.append((top_left, bottom_right))
    
    # 如果有重叠的匹配区域，进行去重处理
    if positions:
        positions = non_max_suppression(positions, overlap_threshold=0.5)
    
    return positions


def non_max_suppression(boxes, overlap_threshold):
    """
    非极大值抑制，去除重复的匹配框
    
    Args:
        boxes (list): 包含(top_left, bottom_right)坐标元组的列表
        overlap_threshold (float): 重叠阈值
    
    Returns:
        list: 去重后的匹配框列表
    """
    if len(boxes) == 0:
        return []
    
    # 转换为numpy数组
    boxes_array = np.array([(box[0][0], box[0][1], box[1][0], box[1][1]) for box in boxes])
    
    # 计算每个框的面积
    areas = (boxes_array[:, 2] - boxes_array[:, 0] + 1) * (boxes_array[:, 3] - boxes_array[:, 1] + 1)
    
    # 根据匹配分数排序（这里简化处理，按x坐标排序）
    indices = np.argsort(boxes_array[:, 0])
    
    keep = []
    while len(indices) > 0:
        # 保留当前索引
        current = indices[0]
        keep.append(current)
        
        # 计算当前框与其他框的重叠比例
        x1_overlap = np.maximum(boxes_array[current, 0], boxes_array[indices[1:], 0])
        y1_overlap = np.maximum(boxes_array[current, 1], boxes_array[indices[1:], 1])
        x2_overlap = np.minimum(boxes_array[current, 2], boxes_array[indices[1:], 2])
        y2_overlap = np.minimum(boxes_array[current, 3], boxes_array[indices[1:], 3])
        
        # 计算重叠面积
        overlap_width = np.maximum(0, x2_overlap - x1_overlap + 1)
        overlap_height = np.maximum(0, y2_overlap - y1_overlap + 1)
        overlap_area = overlap_width * overlap_height
        
        # 计算重叠比例
        overlap_ratio = overlap_area / (areas[current] + areas[indices[1:]] - overlap_area)
        
        # 保留重叠小于阈值的框
        indices = indices[np.where(overlap_ratio <= overlap_threshold)[0] + 1]
    
    # 转换回原始格式
    result_boxes = []
    for i in keep:
        top_left = (int(boxes_array[i, 0]), int(boxes_array[i, 1]))
        bottom_right = (int(boxes_array[i, 2]), int(boxes_array[i, 3]))
        result_boxes.append((top_left, bottom_right))
    
    return result_boxes


def draw_rectangles_on_screenshot(screenshot_path, positions, output_path):
    """
    在屏幕截图上绘制多个矩形标记匹配位置
    
    Args:
        screenshot_path (str): 屏幕截图路径
        positions (list): 包含(top_left, bottom_right)坐标的元组列表
        output_path (str): 输出图片路径
    """
    screenshot = cv2.imread(screenshot_path)
    
    # 为每个匹配位置绘制矩形
    for (top_left, bottom_right) in positions:
        cv2.rectangle(screenshot, top_left, bottom_right, (0, 0, 255), 2)
    
    cv2.imwrite(output_path, screenshot)


if __name__ == "__main__":
    # 定义文件路径
    template_img_path = r"D:\PycharmProjects\auto-script\img\at.png"
    screenshot_path = r"D:\PycharmProjects\auto-script\screentshot\3.png"
    output_path = r"D:\PycharmProjects\auto-script\screentshot\3_template_matched.png"
    
    try:
        # 查找模板位置
        positions = find_template_in_screenshot(template_img_path, screenshot_path, threshold=0.8)
        
        if positions:
            print(f"找到 {len(positions)} 个匹配位置:")
            for i, (top_left, bottom_right) in enumerate(positions):
                print(f"位置 {i+1}: 左上角 {top_left}, 右下角 {bottom_right}")
            
            # 在截图上绘制矩形并保存
            draw_rectangles_on_screenshot(screenshot_path, positions, output_path)
            print(f"标记结果已保存到: {output_path}")
        else:
            print("未找到匹配位置")
            
    except Exception as e:
        print(f"发生错误: {e}")